import os
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, random_split
from torchmetrics import Accuracy
from torchvision import transforms
from torchvision.datasets import MNIST
import matplotlib.pyplot as plt
PATH_DATASETS = "" # 预设路径
BATCH_SIZE = 1024  # 批量
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

# 下载 MNIST 手写阿拉伯数字 训练资料
train_ds = MNIST(PATH_DATASETS, train=True, download=False,
                 transform=transforms.ToTensor())

# 下载测试资料
test_ds = MNIST(PATH_DATASETS, train=False, download=False,
                 transform=transforms.ToTensor())

# 训练/测试资料的维度
# print(train_ds.data.shape, test_ds.data.shape)
#
# print(train_ds.data[0])

# data = train_ds.data[0].clone()
# # data[data >0] = 1
# # data = data.numpy()
# # text_image=[]
# # for i in range(data.shape[0]):
# #     print(''.join(data[i].astype(str)))
# #     text_image.append(''.join(data[i].astype(str)))
#
# plt.imshow(data, cmap='gray')
# plt.show()
# print(text_image)

model = torch.nn.Sequential(
    torch.nn.Flatten(),
    torch.nn.Linear(28 * 28, 256),
    torch.nn.Dropout(0.2),
    torch.nn.Linear(256,10)
)
epochs = 5
lr = 0.1
train_loader = DataLoader(train_ds,batch_size=600)
optimizer = torch.optim.Adadelta(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
model.train()
loss_list = []
for epoch in range(1,epochs+1):
   for batch_idx, (data, target) in enumerate(train_loader):
       data, target = data.to(device), target.to(device)
       optimizer.zero_grad()
       output = model(data)
       loss = criterion(output, target)
       loss.backward()
       optimizer.step()
       if batch_idx % 10 == 0:
           loss_list.append(loss.item())
           batch = batch_idx *len(data)
           data_count = len(train_loader.dataset)
           percentage = (100. * batch_idx/len(train_loader))
           print(f'Epoch {epoch}:[{batch:5d}/{data_count}] ({percentage:.0f}%)' +
                 f' loss:{loss.item():.4f}')

torch.save(model,'model.pt')